import rtde_control
import rtde_receive
import sys
import threading
import numpy as np
from compensation import GravityCompensation

ip='192.168.68.250'

rtde_r = rtde_receive.RTDEReceiveInterface(ip)
rtde_c = rtde_control.RTDEControlInterface(ip)
compensation = GravityCompensation()
# print(rtde_c.zeroFtSensor())

collect_data_pos = [[-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,np.pi/2,-np.pi],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, -3*np.pi/4],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi/2],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi/4],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,np.pi/2,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,np.pi/4,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,0,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi/4,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-3*np.pi/4,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi/2,-np.pi,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-3*np.pi/4,-np.pi/2,0],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-np.pi,-np.pi/2,0],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -np.pi/4],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -np.pi / 2],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -3*np.pi / 4],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -np.pi, -np.pi / 2, -np.pi],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi/4*3],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi/2],
                        [-3 * np.pi / 2, -np.pi / 2, -np.pi / 2, -5 * np.pi / 4, -np.pi / 2, -np.pi/4],
                        [-3*np.pi/2,-np.pi/2,-np.pi/2,-5*np.pi/4,-np.pi/2,0]]
def print_force_torque():
    while True:
        force = rtde_r.getFtRawWrench()
        q = rtde_r.getActualQ()
        rotation_vector = rtde_c.getForwardKinematics(q,[0,0,0,0,0,0])
        force = compensation.GetForceTorque(force, rotation_vector[3:6])
        print(force)
        # sys.stdout.write('\r')
        # sys.stdout.write(str(force))
        # sys.stdout.flush()

# thread = threading.Thread(target=print_force_torque)
# thread.start()

for i in range(len(collect_data_pos)):
    rtde_c.moveJ(collect_data_pos[i],0.2,0.05)
    force = rtde_r.getFtRawWrench()
    # print(force)
    q = rtde_r.getActualQ()
    rotation_vector = rtde_c.getForwardKinematics(q, [0, 0, 0, 0, 0, 0])
    force = compensation.GetForceTorque(force, rotation_vector[3:6])
    print(force)

